import os
from neo4j import GraphDatabase
import random
from faker import Faker

# 初始化Faker用于生成假数据
fake = Faker()

class Neo4jTestDataGenerator:
    def __init__(self, uri, auth):
        self.driver = GraphDatabase.driver(uri, auth=auth)
        
    def close(self):
        self.driver.close()
        
    def empty_first(self):
        """避免数据重复，先清空旧数据"""
        with self.driver.session() as session:
            session.run("""
            MATCH (p:Person),(c:Company)
            DETACH DELETE p,c
            """)
        
    def create_persons(self, count):
        """创建指定数量的Person节点"""
        with self.driver.session() as session:
            for _ in range(count):
                name = fake.name()
                age = random.randint(18, 65)
                email = fake.email()
                job = fake.job()
                
                session.run(
                    "CREATE (p:Person {name: $name, age: $age, email: $email, job: $job})",
                    name=name, age=age, email=email, job=job
                )
            print(f"已创建 {count} 个Person节点")
            
    def create_companies(self, count):
        """创建指定数量的Company节点"""
        with self.driver.session() as session:
            industries = ["科技", "金融", "教育", "医疗", "零售", "制造", "娱乐"]
            
            for _ in range(count):
                name = fake.company()
                industry = random.choice(industries)
                founded_year = random.randint(1950, 2023)
                employees = random.randint(10, 10000)
                
                session.run(
                    "CREATE (c:Company {name: $name, industry: $industry, founded_year: $founded_year, employees: $employees})",
                    name=name, industry=industry, founded_year=founded_year, employees=employees
                )
            print(f"已创建 {count} 个Company节点")
            
    def create_relationships(self):
        """创建节点之间的关系"""
        with self.driver.session() as session:
            # 创建"工作于"关系  # WHERE rand() < 0.1  10%的概率建立关系
            session.run("""
                MATCH (p:Person), (c:Company)
                WITH p, c
                WHERE rand() < 0.1 
                CREATE (p)-[:WORKS_AT {since: rand()*(2023-2000)+2000}]->(c)
            """)
            
            # 创建"朋友"关系 #  rand() < 0.05  避免重复关系，5%的概率建立朋友关系
            session.run("""
                MATCH (p1:Person), (p2:Person)
                WHERE elementId(p1) < elementId(p2) AND rand() < 0.05
                CREATE (p1)-[:FRIENDS_WITH {since: rand()*(2023-2010)+2010}]->(p2)
            """)
            
            # 创建"毕业于"关系（学校可以视为特殊的Company）
            session.run("""
                MATCH (p:Person), (c:Company)
                WHERE c.industry = "教育" AND rand() < 0.15
                CREATE (p)-[:GRADUATED_FROM {year: rand()*(2023-2000)+2000}]->(c)
            """)
            
            print("已创建关系")
            
    def create_test_data(self, num_persons=50, num_companies=10):
        """创建完整的测试数据集"""
        self.empty_first()
        self.create_persons(num_persons)
        self.create_companies(num_companies)
        self.create_relationships()
        print("测试数据生成完成")

if __name__ == "__main__":
    # export NEO4J_URI="bolt://localhost:7687"
    # export NEO4J_USER_NAME="neo4j"
    # export NEO4J_PASSWORD="password"
    # Neo4j连接配置
    neo4j_uri = os.environ.get('NEO4J_URI')
    neo4j_user_name = os.environ.get('NEO4J_USER_NAME')
    neo4j_password = os.environ.get('NEO4J_PASSWORD')
    # 创建生成器实例
    generator = Neo4jTestDataGenerator(neo4j_uri, (neo4j_user_name, neo4j_password))
    
    try:
        # 生成测试数据：50个人，10个公司
        generator.create_test_data(50, 10)
    finally:
        # 关闭连接
        generator.close()
